package com.shujia.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo5ImageModelTrain {
  def main(args: Array[String]): Unit = {

    /**
     * 1、加载特征工程处理过后的数据
     */
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[*]")
      .appName("Demo4ReadImage")
      .config("spark.sql.shuffle.partitions", 4)
      .getOrCreate()

    val imagesData: DataFrame = spark
      .read
      .format("libsvm")
      .load("spark/data/images")
      .coalesce(20)

    /**
     * 2、切分训练集、测试集
     */

    val arr: Array[Dataset[Row]] = imagesData.randomSplit(Array(0.7, 0.3))
    val trainData: Dataset[Row] = arr(0)
    val testData: Dataset[Row] = arr(1)
    //    println(trainData.count())
    //    println(testData.count())

    //    trainData.show(10,false)
    //    testData.show(10,false)

    /**
     * 3、选择合适的模型进行训练
     */
    val lr = new LogisticRegression()
      .setMaxIter(10)
      .setFitIntercept(true)

    val model: LogisticRegressionModel = lr.fit(trainData) // 进行模型训练

    /**
     * 4、进行模型评估
     */
    val resDF: DataFrame = model.transform(testData)
    import spark.implicits._
    import org.apache.spark.sql.functions._

    // 将数据中原有的label 同 预测出来的结果 prediction 进行比较 计算模型的准确率

    resDF
      .select(sum(when($"label" === $"prediction", 1).otherwise(0)) / count($"label"))
      .show()

    /**
     * 5、模型评估后符合要求，将模型进行保存，以便下一次应用
     * 一般要求准确率达到95%
     */
    model.save("spark/data/imageModel")

  }

}
